import cv2 as cv
import numpy as np
import sys
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D  # 空间三维画图

#图像unsharp mask 增强
def unsharp_mask(im, k = 1.0, r = 31):
    '''
    im: input image
    k: degree of enhancing
    r: radius
    '''
    blured = cv.GaussianBlur(im, (r,r), 0)
    enhanced = (k + 1) * im - k * blured
    enhanced[enhanced < 0] = 0
    enhanced[enhanced > 255] = 255
    enhanced = enhanced.astype(np.uint8)
    return enhanced

#彩色图自适应二值化
def color2bin(im, blk_size = 41, C = 15):
    '''
    im: input color image
    blk_size:  block size for adaptiveThreshold
    C:   C paramter for adaptiveThreshold
    '''
    gray = cv.cvtColor(im, cv.COLOR_BGR2GRAY)
    bin = cv.adaptiveThreshold(gray, 255, cv.ADAPTIVE_THRESH_MEAN_C, cv.THRESH_BINARY_INV, blk_size, C)
    return bin

#二值化图像膨胀操作
def dialtion(im, k=1):
    element = cv.getStructuringElement(cv.MORPH_RECT, (2*k+1,2*k+1))
    im_dia = cv.dilate(im, element)
    return im_dia

if __name__ == '__main__':
    fn = ""
    if len(sys.argv) < 2:
        fn = '0-input.png'
    else:
        fn = sys.argv[1]
    
    im = cv.imread(fn)
    sharped = unsharp_mask(im,k=1.5)
    #cv.imwrite('1-sharped.png', sharped)

    bin = color2bin(sharped)
    #cv.imwrite('2-bin.png', 255 - bin)
    
    dia_bin = dialtion(bin, 1)
    #cv.imwrite('3-dialted.png', dia_bin)

    edge = np.bitwise_xor(bin, dia_bin)
    #cv.imwrite('4-edge.png', edge)
    
    #将RGB转为HSV，减小像素通道之间的相关性，提高聚类算法的精度
    im_hsv = cv.cvtColor(im, cv.COLOR_BGR2HSV_FULL)

    #提取参数聚类的像素
    pixels = im_hsv[dia_bin > 0]

    # 绘制像素散点图
    fig = plt.figure()
    ax = ax = fig.add_subplot(projection='3d')#Axes3D(fig)
    im_rgb = cv.cvtColor(im, cv.COLOR_BGR2RGB)
    ax.scatter(pixels[:,0], pixels[:,1], pixels[:,2],marker = '.',c = im_rgb[dia_bin > 0]/255.0)
    ax.set_ylabel('S', fontdict={'size': 15, 'color': 'red'})
    ax.set_xlabel('H', fontdict={'size': 15, 'color': 'red'})
    ax.set_zlabel('V', fontdict={'size': 15, 'color': 'red'})
    plt.show()


